from typing import List, Set, Dict, Optional
from collections import defaultdict
import math

from chernc.treediff.gumtree.matchers.mapping import Mapping
from chernc.treediff.gumtree.matchers.mapping_store import MappingStore
from chernc.treediff.gumtree.matchers.similarity_metrics import SimilarityMetrics
from chernc.treediff.gumtree.tree.tree import TreeLike
from chernc.treediff.gumtree.utils.sequence_algorithms import SequenceAlgorithms

class MappingComparators:
    class FullMappingComparator:
        def __init__(self, ms: 'MappingStore'):
            self.siblings_comparator = MappingComparators.SiblingsSimilarityMappingComparator(ms)
            self.parents_comparator = MappingComparators.ParentsSimilarityMappingComparator()
            self.parents_position_comparator = MappingComparators.PositionInParentsSimilarityMappingComparator()
            self.textual_position_comparator = MappingComparators.TextualPositionDistanceMappingComparator()
            self.position_comparator = MappingComparators.AbsolutePositionDistanceMappingComparator()

        def __call__(self, m1: 'Mapping', m2: 'Mapping') -> int:
            # compare with matched siblings similarity
            result = self.siblings_comparator(m1, m2)
            if result != 0:
                return result

            # compare with ancestors similarity
            result = self.parents_comparator(m1, m2)
            if result != 0:
                return result

            # compare with relative position similarity
            result = self.parents_position_comparator(m1, m2)
            if result != 0:
                return result

            # compare with relative pos
            result = self.textual_position_comparator(m1, m2)
            if result != 0:
                return result

            # compare with absolute pos
            return self.position_comparator(m1, m2)

    class SiblingsSimilarityMappingComparator:
        def __init__(self, ms: 'MappingStore'):
            self.ms = ms
            self.src_descendants: Dict[TreeLike, Set[TreeLike]] = defaultdict(set)
            self.dst_descendants: Dict[TreeLike, Set[TreeLike]] = defaultdict(set)
            self.cached_similarities: Dict['Mapping', float] = {}

        def __call__(self, m1: 'Mapping', m2: 'Mapping') -> int:
            if m1.src.get_parent() == m2.src.get_parent() and m1.dst.get_parent() == m2.dst.get_parent():
                return 0

            if m1 not in self.cached_similarities:
                self.cached_similarities[m1] = SimilarityMetrics.dice_coefficient(
                    self.common_descendants_nb(m1.src.get_parent(), m1.dst.get_parent()),
                    len(self.src_descendants[m1.src.get_parent()]),
                    len(self.dst_descendants[m1.dst.get_parent()])
                )

            if m2 not in self.cached_similarities:
                self.cached_similarities[m2] = SimilarityMetrics.dice_coefficient(
                    self.common_descendants_nb(m2.src.get_parent(), m2.dst.get_parent()),
                    len(self.src_descendants[m2.src.get_parent()]),
                    len(self.dst_descendants[m2.dst.get_parent()])
                )

            return -1 if self.cached_similarities[m1] < self.cached_similarities[m2] else 1

        def common_descendants_nb(self, src: TreeLike, dst: TreeLike) -> int:
            self.src_descendants[src] = set(src.get_descendants())
            self.dst_descendants[dst] = set(dst.get_descendants())

            common = 0
            for t in self.src_descendants[src]:
                m = self.ms.get_dst_for_src(t)
                if m is not None and m in self.dst_descendants[dst]:
                    common += 1

            return common

    class ParentsSimilarityMappingComparator:
        def __init__(self):
            self.src_ancestors: Dict[TreeLike, List[TreeLike]] = defaultdict(list)
            self.dst_ancestors: Dict[TreeLike, List[TreeLike]] = defaultdict(list)
            self.cached_similarities: Dict['Mapping', float] = {}

        def __call__(self, m1: 'Mapping', m2: 'Mapping') -> int:
            if m1.src.get_parent() == m2.src.get_parent() and m1.dst.get_parent() == m2.dst.get_parent():
                return 0

            self.src_ancestors[m1.src] = m1.src.get_parents()
            self.dst_ancestors[m1.dst] = m1.dst.get_parents()
            self.src_ancestors[m2.src] = m2.src.get_parents()
            self.dst_ancestors[m2.dst] = m2.dst.get_parents()

            if m1 not in self.cached_similarities:
                m1_sim = SimilarityMetrics.dice_coefficient(
                    self.common_parents_nb(m1.src, m1.dst),
                    len(self.src_ancestors[m1.src]),
                    len(self.dst_ancestors[m1.dst])
                )
                self.cached_similarities[m1] = m1_sim

            if m2 not in self.cached_similarities:
                m2_sim = SimilarityMetrics.dice_coefficient(
                    self.common_parents_nb(m2.src, m2.dst),
                    len(self.src_ancestors[m2.src]),
                    len(self.dst_ancestors[m2.dst])
                )
                self.cached_similarities[m2] = m2_sim

            return -1 if self.cached_similarities[m1] < self.cached_similarities[m2] else 1

        def common_parents_nb(self, src: TreeLike, dst: TreeLike) -> int:
            return len(SequenceAlgorithms.longest_common_subsequence_with_type(
                self.src_ancestors[src],
                self.dst_ancestors[dst]
            ))

    class PositionInParentsSimilarityMappingComparator:
        def __call__(self, m1: 'Mapping', m2: 'Mapping') -> int:
            m1_distance = self.distance(m1)
            m2_distance = self.distance(m2)
            return -1 if m1_distance < m2_distance else 1

        def distance(self, m: 'Mapping') -> float:
            pos_vector1 = self.pos_vector(m.src)
            pos_vector2 = self.pos_vector(m.dst)
            sum_val = 0
            for i in range(min(len(pos_vector1), len(pos_vector2))):
                sum_val += (pos_vector1[i] - pos_vector2[i]) ** 2
            return math.sqrt(sum_val)

        def pos_vector(self, src: TreeLike) -> List[float]:
            pos_vector = []
            current = src
            while current is not None and current.get_parent() is not None:
                parent = current.get_parent()
                pos = parent.get_child_position(current)
                pos_vector.append(pos)
                current = parent
            return pos_vector

    class TextualPositionDistanceMappingComparator:
        def __call__(self, m1: 'Mapping', m2: 'Mapping') -> int:
            m1_pos_dist = self.textual_position_distance(m1.src, m1.dst)
            m2_pos_dist = self.textual_position_distance(m2.src, m2.dst)
            return -1 if m1_pos_dist < m2_pos_dist else 1

        def textual_position_distance(self, src: TreeLike, dst: TreeLike) -> int:
            return abs(src.get_pos() - dst.get_pos()) + abs(src.get_end_pos() - dst.get_end_pos())

    class AbsolutePositionDistanceMappingComparator:
        def __call__(self, m1: 'Mapping', m2: 'Mapping') -> int:
            m1_pos_dist = self.absolute_position_distance(m1.src, m1.dst)
            m2_pos_dist = self.absolute_position_distance(m2.src, m2.dst)
            return -1 if m1_pos_dist < m2_pos_dist else 1

        def absolute_position_distance(self, src: TreeLike, dst: TreeLike) -> int:
            return abs(src.get_metrics().position - dst.get_metrics().position)